from LogManager import log_output
log_output('预处理：数据生成', send_status=True)
import AFM.handle_sjtu
import AFM.general_dataset
import AFM.out_dataset
import AFM.load_sequence
from PreRecommender import Recommender as PreRecommender
from torch.cuda import is_available
from tqdm import tqdm
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

hyper_params = {
    'dataset_path': 'bslen15/dataset_out_2.txt',
    'bind_file_path': 'bslen15/',
    'load_from_file': False,
    'device': 'cuda' if is_available() else 'cpu',
    'batch_size': 200 if is_available() else 64,
    'epochs': 5 if is_available() else 2,
    'embed_dim': 64,
    'neg_sample_cnt': 300,
    'epsilon': 0.1,
    'total_cnt': 500,
    'bpr_cnt': 200,
    'icf_cnt': 100,
    'ucf_cnt': 100,
    'hot_cnt': 1000
}

def start_preprocess():
    # Pre Recommend
    print('Recommender: Start coarse recommendation...')
    print('Recommender: Initializing...')
    log_output('预处理：初始化粗推荐', send_status=True)
    pre_recommender = PreRecommender(hyper_params)
    print('Recommender: Recommend...')
    log_output('预处理：粗推荐', send_status=True)
    result = pre_recommender.recommend()
    # Write to data
    user_to_bind = pre_recommender.pop_recommender.trainer.file_loader.bind_user_dict
    book_to_bind = pre_recommender.pop_recommender.trainer.file_loader.bind_book_dict

    # Hot Recommendation
    hot_book_dict = {}
    interact_list = pre_recommender.pop_recommender.trainer.file_loader.intereaction_list
    for _, history in interact_list.items():
        for interact in history:
            if interact.item_id not in hot_book_dict:
                hot_book_dict[interact.item_id] = 0

            hot_book_dict[interact.item_id] += 1

    hot_book_list = sorted(hot_book_dict.items(),
                           key=lambda x: x[1], reverse=True)[:hyper_params['hot_cnt']]

    print('Recommender: Writing to file')
    log_output('预处理：写入文件', send_status=True)
    with open(hyper_params['bind_file_path']+'coarse_result.txt', 'wt') as f:
        for user_id, rec_result in result.items():
            bind_user_id = user_to_bind[user_id]
            f.write(f'{bind_user_id};')
            for rec in rec_result:
                f.write(f'{book_to_bind[rec]} ')
            f.write('\n')

    with open(hyper_params['bind_file_path']+'hot_book.txt', 'wt') as f:
        for k, v in hot_book_list:
            f.write(f'{k};{v}\n')

    print('Recommender: Finished coarse recommendation.')


if __name__ == '__main__':
    start_preprocess()
